import os
import random
from PIL import Image, ImageEnhance


def augment_images(input_dir, output_dir, augment_count=4):
    """
    对少量图片进行数据增强，避免非90度旋转。
    :param input_dir: 原始图片路径
    :param output_dir: 增强后图片保存路径
    :param augment_count: 每张图片生成的增强图片数量
    """
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    for img_name in os.listdir(input_dir):
        img_path = os.path.join(input_dir, img_name)
        img = Image.open(img_path).convert("RGB")  # 确保图片为RGB格式

        for i in range(augment_count):
            # 创建增强图片的副本
            augmented_img = img.copy()

            # 随机翻转
            if random.choice([True, False]):
                augmented_img = augmented_img.transpose(Image.FLIP_LEFT_RIGHT)
            if random.choice([True, False]):
                augmented_img = augmented_img.transpose(Image.FLIP_TOP_BOTTOM)

            # 调整亮度和对比度
            enhancer = ImageEnhance.Brightness(augmented_img)
            augmented_img = enhancer.enhance(random.uniform(0.8, 1.2))  # 亮度
            enhancer = ImageEnhance.Contrast(augmented_img)
            augmented_img = enhancer.enhance(random.uniform(0.8, 1.2))  # 对比度

            # 保存增强图片
            aug_name = f"aug_{i}_{img_name}"
            augmented_img.save(os.path.join(output_dir, aug_name))

        print(f"Augmented {img_name} -> {augment_count} new images")


# 原始缺陷图片路径和输出路径
input_dir = "rawA"  # 原始缺陷图片文件夹
output_dir = "rawA_enhance"  # 增强后的图片保存文件夹
augment_images(input_dir, output_dir, augment_count=3)
